"""
车牌识别模块
使用PaddleOCR进行车牌字符识别
"""

import numpy as np
import cv2
from typing import List, Tuple, Dict, Any, Optional
from paddleocr import PaddleOCR
import re


class PlateRecognizer:
    """
    车牌识别器类
    
    功能：
    - 使用PaddleOCR识别车牌字符
    - 支持中文车牌识别
    - 自动纠正识别结果
    """
    
    # 类级别的OCR实例，实现单例模式
    _ocr_instance = None
    _ocr_initialized = False
    
    def __init__(self, config: Optional[dict] = None):
        """
        初始化车牌识别器
        
        参数:
            config (dict, optional): 配置字典
        """
        self.config = config or {}
        self.use_angle_cls = self.config.get('use_angle_cls', True)
        self.lang = self.config.get('lang', 'ch')
        self.rec_threshold = self.config.get('rec_threshold', 0.5)
        
        # 使用单例模式初始化OCR引擎
        if not PlateRecognizer._ocr_initialized:
            PlateRecognizer._ocr_instance = self._init_ocr()
            PlateRecognizer._ocr_initialized = True
        
        self.ocr = PlateRecognizer._ocr_instance
        
        # 中国车牌字符集
        self.chinese_provinces = [
            '京', '津', '沪', '渝', '冀', '豫', '云', '辽', '黑', '湘',
            '皖', '鲁', '新', '苏', '浙', '赣', '鄂', '桂', '甘', '晋',
            '蒙', '陕', '吉', '闽', '贵', '粤', '青', '藏', '川', '宁',
            '琼', '使', '领'
        ]
        
        # 车牌字符映射（用于纠正常见错误）
        self.char_mapping = {
            'O': '0', 'o': '0',
            'I': '1', 'l': '1',
            'Z': '2', 'z': '2',
            'S': '5', 's': '5',
            'B': '8', 'b': '8',
            'G': '6', 'g': '6',
            'Q': '0', 'q': '0',
            'D': '0',
        }
        
    def _init_ocr(self) -> PaddleOCR:
        """
        初始化PaddleOCR引擎
        
        返回:
            PaddleOCR: OCR引擎实例
        """
        try:
            # 使用基本配置，只保留确定支持的参数
            ocr = PaddleOCR(
                use_angle_cls=self.use_angle_cls,
                lang=self.lang,
            )
            print("PaddleOCR初始化成功，使用系统默认模型目录")
            return ocr
        except Exception as e:
            print(f"PaddleOCR初始化失败: {e}")
            # 尝试最简配置
            try:
                ocr = PaddleOCR(lang='ch')
                print("PaddleOCR初始化成功（最简配置），使用系统默认模型目录")
                return ocr
            except Exception as e2:
                print(f"PaddleOCR最简配置也失败: {e2}")
                raise e2
            
    def recognize(self, plate_image: np.ndarray) -> Dict[str, Any]:
        """
        识别车牌图像中的字符
        
        参数:
            plate_image (np.ndarray): 车牌图像
            
        返回:
            Dict[str, Any]: 识别结果，包含:
                - text: 识别的车牌号
                - confidence: 置信度
                - raw_text: 原始识别文本
        """
        # 预处理图像
        processed_image = self._preprocess_plate_image(plate_image)
        
        # 执行OCR识别
        try:
            # PaddleOCR 新版本API - 不再使用cls参数
            result = self.ocr.ocr(processed_image)
            
            if result and result[0]:
                # 提取识别文本和置信度
                text = result[0][0][1][0]
                confidence = result[0][0][1][1]
                
                # 后处理识别结果
                corrected_text = self._postprocess_plate_text(text)
                
                return {
                    'text': corrected_text,
                    'confidence': confidence,
                    'raw_text': text
                }
            else:
                return {
                    'text': '',
                    'confidence': 0.0,
                    'raw_text': ''
                }
        except Exception as e:
            print(f"识别错误: {e}")
            return {
                'text': '',
                'confidence': 0.0,
                'raw_text': ''
            }
            
    def recognize_batch(self, plate_images: List[np.ndarray]) -> List[Dict[str, Any]]:
        """
        批量识别车牌图像
        
        参数:
            plate_images (List[np.ndarray]): 车牌图像列表
            
        返回:
            List[Dict[str, Any]]: 识别结果列表
        """
        results = []
        for plate_image in plate_images:
            result = self.recognize(plate_image)
            results.append(result)
        return results
        
    def _preprocess_plate_image(self, image: np.ndarray) -> np.ndarray:
        """
        预处理车牌图像以提高识别率
        
        参数:
            image (np.ndarray): 原始车牌图像
            
        返回:
            np.ndarray: 预处理后的图像
        """
        # 确保是灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        # 调整大小到标准尺寸
        height, width = gray.shape
        if width < 200 or height < 60:
            scale = max(200 / width, 60 / height)
            new_width = int(width * scale)
            new_height = int(height * scale)
            gray = cv2.resize(gray, (new_width, new_height), 
                            interpolation=cv2.INTER_CUBIC)
        
        # 增强对比度
        gray = cv2.equalizeHist(gray)
        
        # 二值化
        _, binary = cv2.threshold(gray, 0, 255, 
                                cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # 去噪
        binary = cv2.medianBlur(binary, 3)
        
        # 转回三通道（PaddleOCR需要）
        result = cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
        
        return result
        
    def _postprocess_plate_text(self, text: str) -> str:
        """
        后处理识别的车牌文本
        
        参数:
            text (str): 原始识别文本
            
        返回:
            str: 处理后的车牌号
        """
        # 移除空格和特殊字符
        text = re.sub(r'[^\w\u4e00-\u9fa5]', '', text)
        text = text.upper()
        
        # 纠正常见的识别错误
        corrected = ''
        for i, char in enumerate(text):
            if i == 0:  # 第一位应该是省份简称
                if char in self.chinese_provinces:
                    corrected += char
                else:
                    # 尝试找到最相似的省份
                    corrected += self._find_similar_province(char)
            elif i == 1:  # 第二位应该是字母
                if char.isalpha():
                    corrected += char
                else:
                    # 尝试纠正为字母
                    corrected += self.char_mapping.get(char, 'A')
            else:  # 其余位置是字母或数字
                if char in self.char_mapping and i > 1:
                    corrected += self.char_mapping[char]
                else:
                    corrected += char
                    
        # 确保车牌长度正确（普通车牌7位，新能源8位）
        if len(corrected) == 7 or len(corrected) == 8:
            return corrected
        elif len(corrected) > 8:
            return corrected[:8]
        else:
            return corrected
            
    def _find_similar_province(self, char: str) -> str:
        """
        找到最相似的省份简称
        
        参数:
            char (str): 输入字符
            
        返回:
            str: 最相似的省份简称
        """
        # 简单的相似度匹配，实际项目中可以使用更复杂的算法
        if char in self.chinese_provinces:
            return char
        else:
            # 返回默认省份
            return '京'
            
    def validate_plate_format(self, plate_text: str) -> bool:
        """
        验证车牌格式是否正确
        
        参数:
            plate_text (str): 车牌文本
            
        返回:
            bool: 是否为有效格式
        """
        # 普通车牌正则
        normal_pattern = r'^[' + ''.join(self.chinese_provinces) + r'][A-Z][A-Z0-9]{5}$'
        # 新能源车牌正则
        new_energy_pattern = r'^[' + ''.join(self.chinese_provinces) + r'][A-Z]([A-Z0-9]{5}|[DF][A-Z0-9]{5})$'
        
        return bool(re.match(normal_pattern, plate_text) or 
                   re.match(new_energy_pattern, plate_text)) 